# -*- coding:utf-8 -*-

import numpy as np


class FactoryOptimization:
    def __init__(self, months=12, alpha=0.1, gamma=0.9, epsilon=0.1):
        self.months = months  # 每年12个月
        self.alpha = alpha  # 学习率
        self.gamma = gamma  # 折扣因子
        self.epsilon = epsilon  # 探索率

        # 初始化每个工厂的状态空间
        self.states_A = self.init_states_A()
        self.states_B = self.init_states_B()
        self.states_C = self.init_states_C()

        self.actions = self.init_actions()
        self.q_table_A = np.zeros((len(self.states_A), len(self.actions)))
        self.q_table_B = np.zeros((len(self.states_B), len(self.actions)))
        self.q_table_C = np.zeros((len(self.states_C), len(self.actions)))

    def init_states_A(self):
        states = []
        for x11 in range(100, 201, 50):
            for y11 in np.arange(0.8, 1.1, 0.1):
                for c11 in range(10, 31, 10):
                    states.append((x11, y11, c11))
        return states

    def init_states_B(self):
        states = []
        for x21 in range(100, 201, 50):
            for y21 in np.arange(0.8, 1.1, 0.1):
                for c21 in range(10, 31, 10):
                    states.append((x21, y21, c21))
        return states

    def init_states_C(self):
        states = []
        for x31 in range(100, 201, 50):
            for y31 in np.arange(0.8, 1.1, 0.1):
                for c31 in range(10, 31, 10):
                    states.append((x31, y31, c31))
        return states

    def init_actions(self):
        return ['increase_x', 'decrease_x', 'increase_y', 'decrease_y']

    # def reward(self, state, action, factory):
    #     x, y, c = list(state)  # 将状态解包为变量
    #
    #     # 根据动作调整状态变量
    #     if action == 'increase_x':
    #         x += 10  # 增加投入量
    #     elif action == 'decrease_x':
    #         x = max(0, x - 10)  # 减少投入量，确保不为负
    #     elif action == 'increase_y':
    #         y = min(1.0, y + 0.1)  # 增加合格率，确保不超过1.0
    #     elif action == 'decrease_y':
    #         y = max(0.0, y - 0.1)  # 减少合格率，确保不为负
    #
    #     # 重新计算利润
    #     if factory == 'A':
    #         p1, p2, p4 = 100, 120, 150
    #         q1, q2, q4 = x, x, x  # 假设产量与投入量成正比
    #         profit_g1 = (p1 - c) * q1 * y
    #         profit_g2 = (p2 - c) * q2 * y
    #         profit_g4 = (p4 - c) * q4 * y
    #         total_profit = profit_g1 + profit_g2 + profit_g4
    #
    #     elif factory == 'B':
    #         p1, p3, p5 = 100, 130, 160
    #         q1, q3, q5 = x, x, x
    #         profit_g1 = (p1 - c) * q1 * y
    #         profit_g3 = (p3 - c) * q3 * y
    #         profit_g5 = (p5 - c) * q5 * y
    #         total_profit = profit_g1 + profit_g3 + profit_g5
    #
    #     elif factory == 'C':
    #         p3, p4 = 130, 150
    #         q3, q4 = x, x
    #         profit_g3 = (p3 - c) * q3 * y
    #         profit_g4 = (p4 - c) * q4 * y
    #         total_profit = profit_g3 + profit_g4
    #
    #     return total_profit

    def reward(self, state, action, factory):
        x, y, c = list(state)  # 将状态解包为变量

        # 根据动作调整状态变量
        if action == 'increase_x':
            x += 10  # 增加投入量
            y *= 0.95  # 假设增加投入量导致合格率下降5%
            c += 5 * (x / 100) ** 2  # 非线性增加成本，假设与投入量的平方成正比
        elif action == 'decrease_x':
            x = max(0, x - 10)  # 减少投入量
            y *= 1.05  # 假设减少投入量导致合格率提高5%
            c -= 3  # 假设成本减少
            c = max(0, c)  # 确保成本不为负
        elif action == 'increase_y':
            y = min(1.0, y + 0.1)  # 增加合格率
            c += 2  # 假设合格率提高导致成本增加
        elif action == 'decrease_y':
            y = max(0.0, y - 0.1)  # 减少合格率
            c -= 1  # 假设合格率降低导致成本减少
            c = max(0, c)  # 确保成本不为负

        # 重新计算利润
        if factory == 'A':
            p1, p2, p4 = 100, 120, 150
            q1, q2, q4 = x, x, x  # 假设产量与投入量成正比
            profit_g1 = (p1 - c) * q1 * y
            profit_g2 = (p2 - c) * q2 * y
            profit_g4 = (p4 - c) * q4 * y
            total_profit = profit_g1 + profit_g2 + profit_g4

        elif factory == 'B':
            p1, p3, p5 = 100, 130, 160
            q1, q3, q5 = x, x, x
            profit_g1 = (p1 - c) * q1 * y
            profit_g3 = (p3 - c) * q3 * y
            profit_g5 = (p5 - c) * q5 * y
            total_profit = profit_g1 + profit_g3 + profit_g5

        elif factory == 'C':
            p3, p4 = 130, 150
            q3, q4 = x, x
            profit_g3 = (p3 - c) * q3 * y
            profit_g4 = (p4 - c) * q4 * y
            total_profit = profit_g3 + profit_g4

        return total_profit

    def choose_action(self, state, q_table):
        if np.random.rand() < self.epsilon:
            return np.random.choice(self.actions)
        else:
            if q_table is self.q_table_A:
                state_index = self.states_A.index(state)
            elif q_table is self.q_table_B:
                state_index = self.states_B.index(state)
            else:
                state_index = self.states_C.index(state)
            return self.actions[np.argmax(q_table[state_index])]

    def update_q_table(self, state, action, reward, next_state, q_table):
        if q_table is self.q_table_A:
            state_index = self.states_A.index(state)
            next_state_index = self.states_A.index(next_state)
        elif q_table is self.q_table_B:
            state_index = self.states_B.index(state)
            next_state_index = self.states_B.index(next_state)
        else:
            state_index = self.states_C.index(state)
            next_state_index = self.states_C.index(next_state)

        action_index = self.actions.index(action)
        q_predict = q_table[state_index, action_index]
        q_target = reward + self.gamma * np.max(q_table[next_state_index])
        q_table[state_index, action_index] += self.alpha * (q_target - q_predict)

    def train(self, epochs=1000):
        total_profit_per_month = []

        for _ in range(epochs):
            monthly_profit = 0

            for month in range(self.months):
                # 训练A工厂
                state_A = self.states_A[np.random.choice(len(self.states_A))]
                action_A = self.choose_action(state_A, self.q_table_A)
                next_state_A = self.states_A[np.random.choice(len(self.states_A))]
                reward_A = self.reward(state_A, action_A, 'A')
                self.update_q_table(state_A, action_A, reward_A, next_state_A, self.q_table_A)

                # 训练B工厂
                state_B = self.states_B[np.random.choice(len(self.states_B))]
                action_B = self.choose_action(state_B, self.q_table_B)
                next_state_B = self.states_B[np.random.choice(len(self.states_B))]
                reward_B = self.reward(state_B, action_B, 'B')
                self.update_q_table(state_B, action_B, reward_B, next_state_B, self.q_table_B)

                # 训练C工厂
                state_C = self.states_C[np.random.choice(len(self.states_C))]
                action_C = self.choose_action(state_C, self.q_table_C)
                next_state_C = self.states_C[np.random.choice(len(self.states_C))]
                reward_C = self.reward(state_C, action_C, 'C')
                self.update_q_table(state_C, action_C, reward_C, next_state_C, self.q_table_C)

                # 每个月的总利润
                monthly_profit += reward_A + reward_B + reward_C

            total_profit_per_month.append(monthly_profit)

        return total_profit_per_month

    def get_optimal_solution(self):
        optimal_policy_A = []
        optimal_policy_B = []
        optimal_policy_C = []

        for state in self.states_A:
            action = self.choose_action(state, self.q_table_A)
            optimal_policy_A.append((state, action))

        for state in self.states_B:
            action = self.choose_action(state, self.q_table_B)
            optimal_policy_B.append((state, action))

        for state in self.states_C:
            action = self.choose_action(state, self.q_table_C)
            optimal_policy_C.append((state, action))

        return optimal_policy_A, optimal_policy_B, optimal_policy_C


# 实例化优化类并训练模型
optimizer = FactoryOptimization()
monthly_profits = optimizer.train()
optimal_solution_A, optimal_solution_B, optimal_solution_C = optimizer.get_optimal_solution()

# 计算一年的总利润
total_annual_profit = sum(monthly_profits)

# 输出每个月的利润和一年的总利润
print("每个月的总利润:")
for month, profit in enumerate(monthly_profits):
    print(f"第{month + 1}个月: {profit}")

print(f"\n一年的总利润: {total_annual_profit}")

# 输出每个月A、B、C工厂的最优策略
print("A工厂最优策略:")
for state, action in optimal_solution_A:
    print(f"State: {state}, Optimal Action: {action}")

print("\nB工厂最优策略:")
for state, action in optimal_solution_B:
    print(f"State: {state}, Optimal Action: {action}")

print("\nC工厂最优策略:")
for state, action in optimal_solution_C:
    print(f"State: {state}, Optimal Action: {action}")
